package cn.iotab.iotkit.transport.api.access.handler.session;

import lombok.extern.slf4j.Slf4j;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
public class InMemorySessionStorage implements SessionStorage {

    /**
     * productKey --> deviceId --> session
     */
    private final Map<String, Map<String, Session>> productDeviceSessions = new ConcurrentHashMap<>();
    /**
     * sessionId --> session
     */
    private final Map<String, Session> sessions = new ConcurrentHashMap<>();

    private final long sessionTimeout;

    public InMemorySessionStorage(long sessionTimeout) {
        this.sessionTimeout = sessionTimeout;
    }

    private Session checkSessionTimeout(Session session) {
        if (session == null) return null;
        if (System.currentTimeMillis() - session.getLastActivityTime() > sessionTimeout) {
            removeSession(session.getSessionId());
            return null;
        }
        return session;
    }

    @Override
    public Session getBySessionId(String sessionId) {
        if (sessionId == null) return null;
        Session session = sessions.get(sessionId);
        return checkSessionTimeout(session);
    }

    @Override
    public Session getByDevice(String productKey, String deviceId) {
        if (productKey == null || deviceId == null) return null;
        Map<String, Session> deviceSession = productDeviceSessions.get(productKey);
        if (deviceSession == null) return null;
        Session session = deviceSession.get(deviceId);
        return checkSessionTimeout(session);
    }

    @Override
    public synchronized void saveSession(Session session) {
        String sessionId = session.getSessionId();
        if (sessionId == null) {
            log.warn("session id is null: {}", session);
            return;
        }
        sessions.put(sessionId, session);
        if (session.getProductKey() != null || session.getDeviceId() != null) {
            Map<String, Session> map = productDeviceSessions.computeIfAbsent(session.getProductKey(), k -> new HashMap<>());
            map.put(session.getDeviceId(), session);
        }
    }

    @Override
    public synchronized Session removeSession(String sessionId) {
        Session session = doRemoveSession(sessionId);
        if (session != null && session.getProductKey() != null && session.getDeviceId() != null) {
            doRemoveSession(session.getProductKey(), session.getDeviceId());
        }
        return session;
    }

    private Session doRemoveSession(String sessionId) {
        return sessions.remove(sessionId);
    }

    private Session doRemoveSession(String productKey, String deviceId) {
        Map<String, Session> map = productDeviceSessions.get(productKey);
        if (map != null) {
            Session session = map.remove(deviceId);
            if (map.isEmpty()) {
                productDeviceSessions.remove(productKey);
            }
            return session;
        }
        return null;
    }
}
